import pysam
import argparse
import pod5 as p5
import matplotlib.pyplot as plt
from deepsignal3.utils import bam_reader
from deepsignal3.extract_features_pod5 import _group_signals_by_movetable_v2
import numpy as np

def process(readids,bam_index, pod5_dr, out_file):
    seqsig_dict = dict()
    for read_name in readids:
        read = pod5_dr.get_read(read_name)  
        signal = read.signal  
        for bam_read in bam_index.get_alignments(read_name):
            if bam_read.is_supplementary or bam_read.is_secondary:
                continue
            reference_name = bam_read.reference_name
            seq = bam_read.get_forward_sequence().upper()
            read_dict = dict(bam_read.tags)
            mv_table = read_dict["mv"]
            num_trimmed = read_dict["ts"]
            if num_trimmed >= 0:
                # (signal[num_trimmed:] - norm_shift) / norm_scale
                signal_trimmed = signal[num_trimmed:]
            else:
                # (signal[:num_trimmed] - norm_shift) / norm_scale
                signal_trimmed = signal[:num_trimmed]
            signal_group = _group_signals_by_movetable_v2(
                        signal_trimmed, np.asarray(mv_table[1:]), int(mv_table[0]))
            seqsig_dict[read_name] = (seq,signal_group)

    return seqsig_dict

def main():
    # 设置命令行参数解析器
    parser = argparse.ArgumentParser(description="")
    
    # 添加参数
    parser.add_argument("--bam", type=str, required=True)
    parser.add_argument("--pod5", type=str, required=True)
    parser.add_argument("--out", type=str, required=True)
    args = parser.parse_args()

    pod5_dr = p5.DatasetReader(args.pod5, recursive=True)
    #bam_file = pysam.AlignmentFile(args.bam, "rb")
    bam_index = bam_reader.ReadIndexedBam(args.bam)